
#include "non_max_suppression_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"   //GetPlatformInfo类接口
#include <algorithm>


/*
boxes：包围框，张量（浮点）。形状为  [num_batches、spatial_dimension、4] 的输入张量。单框数据格式由 center_point_box 表示。
scores：分数，张量（浮点）。形状为 [num_batches、num_classes、spatial_dimension] 的输入张量
max_output_boxes_per_class（可选）：张量（int64）。表示每个类别每个批次要选择的最大框数的整数。它是一个标量。默认为 0，表示无输出。
iou_threshold（可选）：张量（浮点数）。浮点数表示用于决定框相对于 IOU 是否重叠过多的阈值。它是标量。值范围 [0, 1]。默认为 0
score_threshold（可选）：张量（浮点数）。浮点数表示用于根据分数决定何时删除框的阈值。它是一个标量。
属性：

center_point_box：int（默认为 0）。整数表示框数据的格式。默认值为 0。0 — 框数据以  [y1, x1, y2, x2] 的形式提供，其中 (y1, x1) 和 (y2, x2) 是任何对角线框角对的坐标，坐标可以作为归一化（即位于区间 [0, 1] 内）或绝对值提供。主要用于 TF 模型。1 — 框数据以 [x_center, y_center, width, height] 的形式提供。主要用于 Pytorch 模型。
输出：

selected_indices：张量（int64）。从框张量中选择的索引
[num_selected_indices，3]，选定的索引格式为 [batch_index，class_index，box_index]。*/



namespace optiling {

const uint32_t BLOCK_SIZE = 32;   //32bytes数据对齐
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    NonMaxSuppressionTilingData tiling;

    //num_batches, num_classes, spatial_dim
    auto shape_scores = context->GetInputShape(1)->GetOriginShape();

    int32_t num_batches = shape_scores.GetDim(0);
    int32_t num_classes = shape_scores.GetDim(1);
    int32_t spatial_dim = shape_scores.GetDim(2);

    //传递center_point_box参数
    int center_point_box = *context->GetAttrs()->GetInt(0);


    tiling.set_num_batches(num_batches);
    tiling.set_num_classes(num_classes);
    tiling.set_spatial_dim(spatial_dim);
    tiling.set_center_point_box(center_point_box);


    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    auto shape_score = context->GetInputShape(1);  //->GetOriginShape()
    
    auto shape_max   = context->GetInputShape(2);  //->GetOriginShape()

    gert::Shape* y_shape = context->GetOutputShape(0);

    int32_t x_dimensional = 2; //shape_score->GetDimNum() -1;
    int32_t size =1;
    for(int i = 0; i < x_dimensional; i++)
    {
        size *= shape_score->GetDim(i);
    }
    
    if(shape_max->GetDimNum() != 0)
        size *= shape_max->GetDim(0);
    else
        size = 0;

    //printf("size is %d",size);
    const gert::Shape tmp_shape({size,3});

    *y_shape = tmp_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class NonMaxSuppression : public OpDef {
public:
    explicit NonMaxSuppression(const char* name) : OpDef(name)
    {
        this->Input("boxes")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("scores")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("max_output_boxes_per_class")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_INT32})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("iou_threshold")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("score_threshold")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("selected_indices")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        
        this->Attr("center_point_box").AttrType(OPTIONAL).Int(0);
        
        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend310b");
    }
};

OP_ADD(NonMaxSuppression);
}
